#include "../tools.h"

class Solution {
public:
    
    /**
     * @brief 
     * 
     * 归并         O(m+n) O(m+n)
     * 双指针       O(m+n) O(1)
     * 二分查找     O(log(m+n)) O(1) 
     *      寻找两个有序数组中第 k 小的数，k = (m+n)/2 (m+n)/2+1
     *          min(A[k/2-1], B[k/2-1]) 前最多只有 k-2 个元素比它小
     *          排除了 A[0 ~ k/2-1] 或 B[0 ~ k/2-1]
     *          缩小了查找范围，并且根据缩小值改变 k
     */
    double findKth(vector<int>& nums1, vector<int>& nums2, int k){
        int size1 = nums1.size();
        int size2 = nums2.size();
        int index1 = 0, index2 = 0;
        while(true){
            // 终止条件：有一个数组被排除干净 或 k = 1
            if(index1 == size1) return nums2[index2 + k - 1]; // k 从 1 开始
            if(index2 == size2) return nums1[index1 + k - 1];
            if(k == 1) return min(nums1[index1], nums2[index2]);

            int newindex1 = min(index1 + k / 2 - 1, size1 - 1);
            int newindex2 = min(index2 + k / 2 - 1, size2 - 1);
            if(nums1[newindex1] <= nums2[newindex2]){
                k -= newindex1 - index1 + 1;    // 修改 k
                index1 = newindex1 + 1;         // 缩小范围
            }
            else{
                k -= newindex2 - index2 + 1;
                index2 = newindex2 + 1;
            }
        }
    }

    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int size1 = nums1.size();
        int size2 = nums2.size();
        int size = size1 + size2;
        if(size == 0) return 0.0;
        
        if(size & 1 == 1)
            return findKth(nums1, nums2, size / 2 + 1);
        else
            return (findKth(nums1, nums2, size / 2) + findKth(nums1, nums2, size / 2 + 1)) / 2.0;
    }
};


int main(int argc, char const *argv[])
{
    Solution sol;

    vector<int> nums1 = {1};
    vector<int> nums2 = {2,3,4,5,6};

    // cout << sol.findKth(nums1, nums2, 3);
    cout << sol.findMedianSortedArrays(nums1, nums2) << endl;

    return 0;
}
